"""
缓冲区管理器
实现LRU缓存机制，管理内存中的页面
"""

import time
from typing import Dict, Optional, Tuple
from collections import OrderedDict

class BufferFrame:
    """缓冲区帧"""
    
    def __init__(self, table_name: str, page_id: int, page_data: bytearray):
        """初始化缓冲区帧"""
        self.table_name = table_name
        self.page_id = page_id
        self.page_data = page_data
        self.dirty = False  # 是否被修改
        self.pinned = False  # 是否被固定
        self.last_access = time.time()  # 最后访问时间
    
    def mark_dirty(self):
        """标记为脏页"""
        self.dirty = True
    
    def pin(self):
        """固定页面"""
        self.pinned = True
    
    def unpin(self):
        """取消固定页面"""
        self.pinned = False
    
    def update_access_time(self):
        """更新访问时间"""
        self.last_access = time.time()
    
    def get_key(self) -> Tuple[str, int]:
        """获取页面的唯一标识"""
        return (self.table_name, self.page_id)

class BufferManager:
    """缓冲区管理器"""
    
    def __init__(self, buffer_size: int = 100):
        """初始化缓冲区管理器"""
        self.buffer_size = buffer_size
        self.frames: Dict[Tuple[str, int], BufferFrame] = {}
        self.access_order = OrderedDict()  # 用于LRU算法
        self.storage_manager = None  # 存储管理器引用
    
    def set_storage_manager(self, storage_manager):
        """设置存储管理器"""
        self.storage_manager = storage_manager
    
    def get_page(self, table_name: str, page_id: int) -> Optional[bytearray]:
        """获取页面"""
        key = (table_name, page_id)
        
        # 检查是否在缓冲区中
        if key in self.frames:
            frame = self.frames[key]
            frame.update_access_time()
            self._update_access_order(key)
            return frame.page_data
        
        # 从存储管理器读取页面
        if self.storage_manager:
            page_data = self.storage_manager.read_page(table_name, page_id)
            if page_data:
                return self._load_page(table_name, page_id, page_data)
        
        return None
    
    def put_page(self, table_name: str, page_id: int, page_data: bytearray, dirty: bool = False):
        """将页面放入缓冲区"""
        key = (table_name, page_id)
        
        # 如果页面已存在，更新数据
        if key in self.frames:
            frame = self.frames[key]
            frame.page_data = page_data
            frame.update_access_time()
            if dirty:
                frame.mark_dirty()
            self._update_access_order(key)
            return
        
        # 检查缓冲区是否已满
        if len(self.frames) >= self.buffer_size:
            self._evict_page()
        
        # 创建新的缓冲区帧
        frame = BufferFrame(table_name, page_id, page_data)
        if dirty:
            frame.mark_dirty()
        
        self.frames[key] = frame
        self._update_access_order(key)
    
    def pin_page(self, table_name: str, page_id: int):
        """固定页面，防止被驱逐"""
        key = (table_name, page_id)
        if key in self.frames:
            self.frames[key].pin()
    
    def unpin_page(self, table_name: str, page_id: int):
        """取消固定页面"""
        key = (table_name, page_id)
        if key in self.frames:
            self.frames[key].unpin()
    
    def mark_dirty(self, table_name: str, page_id: int):
        """标记页面为脏页"""
        key = (table_name, page_id)
        if key in self.frames:
            self.frames[key].mark_dirty()
    
    def flush_page(self, table_name: str, page_id: int) -> bool:
        """将页面写回磁盘"""
        key = (table_name, page_id)
        if key not in self.frames:
            return False
        
        frame = self.frames[key]
        if frame.dirty and self.storage_manager:
            success = self.storage_manager.write_page(table_name, page_id, frame.page_data)
            if success:
                frame.dirty = False
            return success
        
        return True
    
    def flush_all(self):
        """将所有脏页写回磁盘"""
        for key, frame in self.frames.items():
            if frame.dirty:
                table_name, page_id = key
                self.flush_page(table_name, page_id)
    
    def _load_page(self, table_name: str, page_id: int, page_data: bytearray) -> bytearray:
        """加载页面到缓冲区"""
        self.put_page(table_name, page_id, page_data)
        return page_data
    
    def _update_access_order(self, key: Tuple[str, int]):
        """更新访问顺序"""
        if key in self.access_order:
            self.access_order.move_to_end(key)
        else:
            self.access_order[key] = True
    
    def _evict_page(self):
        """驱逐页面（LRU算法）"""
        # 找到最久未访问且未被固定的页面
        for key in self.access_order:
            frame = self.frames[key]
            if not frame.pinned:
                # 如果是脏页，先写回磁盘
                if frame.dirty:
                    table_name, page_id = key
                    self.flush_page(table_name, page_id)
                
                # 从缓冲区中移除
                del self.frames[key]
                del self.access_order[key]
                return
        
        # 如果没有可驱逐的页面，抛出异常
        raise Exception("缓冲区已满且所有页面都被固定")
    
    def get_buffer_stats(self) -> Dict[str, int]:
        """获取缓冲区统计信息"""
        dirty_count = sum(1 for frame in self.frames.values() if frame.dirty)
        pinned_count = sum(1 for frame in self.frames.values() if frame.pinned)
        
        return {
            'total_frames': len(self.frames),
            'dirty_frames': dirty_count,
            'pinned_frames': pinned_count,
            'buffer_size': self.buffer_size
        }
    
    def print_buffer_info(self):
        """打印缓冲区信息"""
        stats = self.get_buffer_stats()
        print("缓冲区信息:")
        print(f"  总帧数: {stats['total_frames']}/{stats['buffer_size']}")
        print(f"  脏页数: {stats['dirty_frames']}")
        print(f"  固定页数: {stats['pinned_frames']}")
        
        print("  页面列表:")
        for key, frame in self.frames.items():
            table_name, page_id = key
            status = []
            if frame.dirty:
                status.append("脏")
            if frame.pinned:
                status.append("固定")
            status_str = f"({', '.join(status)})" if status else ""
            print(f"    {table_name}.{page_id} {status_str}")
    
    def clear(self):
        """清空缓冲区"""
        self.flush_all()
        self.frames.clear()
        self.access_order.clear()

# 测试函数
if __name__ == "__main__":
    from storage_manager import StorageManager
    
    # 创建存储管理器和缓冲区管理器
    storage = StorageManager("test_data")
    buffer_manager = BufferManager(5)  # 小缓冲区用于测试
    buffer_manager.set_storage_manager(storage)
    
    # 创建测试表
    schema = {
        'columns': {
            'id': {'type': {'type': 'int'}, 'primary_key': True},
            'name': {'type': {'type': 'varchar', 'length': 50}}
        },
        'primary_key': 'id'
    }
    
    storage.create_table("test", schema)
    
    # 测试缓冲区操作
    print("测试缓冲区管理器...")
    
    # 分配页面
    page_id = storage.allocate_page("test")
    print(f"分配页面: {page_id}")
    
    # 创建测试数据
    test_data = bytearray(4096)
    test_data[0:4] = b'TEST'
    
    # 放入缓冲区
    buffer_manager.put_page("test", page_id, test_data, dirty=True)
    print("页面已放入缓冲区")
    
    # 固定页面
    buffer_manager.pin_page("test", page_id)
    print("页面已固定")
    
    # 打印缓冲区信息
    buffer_manager.print_buffer_info()
    
    # 取消固定并刷新
    buffer_manager.unpin_page("test", page_id)
    buffer_manager.flush_page("test", page_id)
    print("页面已刷新到磁盘")
    
    # 清理
    storage.drop_table("test")
    buffer_manager.clear()
